import os
import torch
from config.dataset import ADE20KConfig

class PSPNetConfig:
    """PSPNet model configuration"""
    
    # Model configuration
    MODEL_NAME = "PSPNet"
    ENCODER_NAME = "resnet34"
    ENCODER_WEIGHTS = None  # Do not use pre-trained weights
    
    # Dataset configuration
    DATASET_CONFIG = ADE20KConfig()
    NUM_CLASSES = DATASET_CONFIG.NUM_CLASSES
    
    # Training configuration
    BATCH_SIZE = 4
    NUM_EPOCHS = 100
    LEARNING_RATE = 0.001
    WEIGHT_DECAY = 1e-4
    
    # Image configuration
    IMAGE_SIZE = (512, 512)
    
    # Path configuration
    CHECKPOINT_DIR = "checkpoints_pspnet"
    LOG_DIR = "logs_pspnet"
    
    # Device configuration
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    
    @classmethod
    def create_directories(cls):
        """Create necessary directories"""
        os.makedirs(cls.CHECKPOINT_DIR, exist_ok=True)
        os.makedirs(cls.LOG_DIR, exist_ok=True)